from transformers import AutoTokenizer, AutoModelForCausalLM

# 加载微调后的模型和分词器
model_path = "./fine_tuned_medical_model"  # 替换为你实际保存微调模型的路径
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path)

# 定义问题
question = "患者，女性，48岁，反复发热15天，咳嗽咳痰10天，行胸部X线检查示右下肺野斑片状模糊影，其内可见有空洞。影像检查的其他补充征象描述有：空洞内可见气液平，空洞外壁边界不清，其周有斑片模糊影，空洞壁较厚且均匀，纵隔窗未见明显淋巴结肿大，进一步行PPD试验结果为阴性，最可能的诊断为？"

# 构建输入文本，模仿训练时的格式
instruction = "请具体分析"
input_text = f"<|start_header_id|>user<|end_header_id|>\n\n{instruction}: {question}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"

# 对输入文本进行分词处理
inputs = tokenizer(input_text, return_tensors="pt")

# 生成回答
outputs = model.generate(
    input_ids=inputs["input_ids"],
    attention_mask=inputs["attention_mask"],
    max_new_tokens=500,  # 控制生成回答的最大长度
    temperature=0.7  # 控制生成文本的随机性，可根据需要调整
)

# 解码生成的回答
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

# 提取回答部分，去除输入部分
answer_start = input_text.rstrip()  # 去除输入文本末尾的空白字符
answer = generated_text[len(answer_start):]

print("模型生成的回答：", answer)
